#include "assembler.hpp"
#include "program.hpp"
#include "opcode.hpp"

#include <boost/foreach.hpp>
#include <boost/fusion/include/adapt_struct.hpp>
#include <boost/optional.hpp>
#include <boost/spirit/include/phoenix.hpp>
#include <boost/spirit/include/qi.hpp>
#include <boost/variant/recursive_variant.hpp>


namespace qi = boost::spirit::qi;
namespace charset = boost::spirit::ascii;
namespace phoenix = boost::phoenix;

struct AddAST;
struct SubAST;
struct DeclarationAST;

typedef std::vector<char> Literal;

struct RandomValue {
	RandomValue(char) {}
	RandomValue() {}
};

typedef boost::variant<
	  int
	, Literal
	, boost::recursive_wrapper<AddAST>
	, boost::recursive_wrapper<SubAST>
	, RandomValue
	, boost::recursive_wrapper<DeclarationAST>
> ConstantAST;

struct AddAST {
	ConstantAST first;
	ConstantAST second;
};

struct SubAST {
	ConstantAST first;
	ConstantAST second;
};

struct InstructionAST {
	ConstantAST id;
	ConstantAST op;
	ConstantAST argA;
	ConstantAST argB;
	ConstantAST ret;
};

struct DeclarationAST {
	Literal dname;
	ConstantAST dvalue;
};

typedef boost::variant<InstructionAST, DeclarationAST> StatementAST;
typedef std::vector<StatementAST> ProgramAST;

BOOST_FUSION_ADAPT_STRUCT(
InstructionAST,
(ConstantAST, id)
(ConstantAST, op)
(ConstantAST, argA)
(ConstantAST, argB)
(ConstantAST, ret)
)

BOOST_FUSION_ADAPT_STRUCT(
AddAST,
(ConstantAST, first)
(ConstantAST, second)
)

BOOST_FUSION_ADAPT_STRUCT(
SubAST,
(ConstantAST, first)
(ConstantAST, second)
)

BOOST_FUSION_ADAPT_STRUCT(
DeclarationAST,
(Literal, dname)
(ConstantAST, dvalue)
)

// Skipper should skip spaces (But leave eol). It should also strip comments (which are denoted by # )
typedef BOOST_TYPEOF( (qi::space - qi::eol) |  (qi::lit('#') >> *(qi::char_ - qi::eol)) ) SkipperType;


template <typename Iterator>
struct ProgramGrammar
	: public qi::grammar<Iterator, ProgramAST(), SkipperType>
{
	ProgramGrammar()
		: ProgramGrammar::base_type(start_)
	{
		add_ %= '(' >> value_ >> '+' >> value_ >> ')';
		sub_ %= '(' >> value_ >> '-' >> value_ >> ')';
		literal_ %= +charset::alpha;
		rand_ = '?';
		value_ %=  rand_ | qi::int_ | literal_ | add_ | sub_ | ('(' >> decl_ >> ')') | ('(' >> value_ >> ')');
		decl_ %= literal_ >> '=' >> value_;
		instruction_ %= value_ >> value_ >> value_ >> value_ >> value_;
		statement_ %= (instruction_ | decl_) >> +(qi::eol); // statements must end in a new line
		start_ %= *qi::eol >> +(statement_);
	}
private:
	qi::rule<Iterator, AddAST(), SkipperType> add_;
	qi::rule<Iterator, SubAST(), SkipperType> sub_;
	qi::rule<Iterator, Literal(), SkipperType> literal_;
	qi::rule<Iterator, RandomValue(), SkipperType> rand_;
	qi::rule<Iterator, ConstantAST(), SkipperType> value_;
	qi::rule<Iterator, DeclarationAST(), SkipperType> decl_;
	qi::rule<Iterator, InstructionAST(), SkipperType> instruction_;
	qi::rule<Iterator, StatementAST(), SkipperType> statement_;
	qi::rule<Iterator, ProgramAST(), SkipperType> start_;
};


struct ConstResolver : public boost::static_visitor< boost::optional<int> > {

	boost::optional<int> operator()(int v) const {
		return v;
	}

	boost::optional<int> operator()(RandomValue) const {
		return ::rand();
	}

	boost::optional<int> operator()(Literal const& v) const {

		std::map<Literal, int>::const_iterator it(lookup_.find(v));

		if (it != lookup_.end())
			return it->second;
		else
			return boost::optional<int>();
	}

	boost::optional<int> operator()(AddAST const& v) const {

		boost::optional<int> fst(boost::apply_visitor(*this, v.first));

		if (fst.is_initialized()) {
			boost::optional<int> snd(boost::apply_visitor(*this, v.second));
			if (snd.is_initialized()) {
				return *fst + *snd;
			}
		}

		return boost::optional<int>();
	}

	boost::optional<int> operator()(SubAST const& v) const {

		boost::optional<int> fst(boost::apply_visitor(*this, v.first));

		if (fst.is_initialized()) {
			boost::optional<int> snd(boost::apply_visitor(*this, v.second));
			if (snd.is_initialized()) {
				return *fst - *snd;
			}
		}

		return boost::optional<int>();
	}

	boost::optional<int> operator()(DeclarationAST const& v) const {
		return (*this)(v.dname);
	}


	void add_constant(std::string const& k, int v) {
		lookup_.insert(std::make_pair(Literal(k.begin(), k.end()), v));
	}

	void add_constant(Literal const& k, int v) {
		lookup_.insert(std::make_pair(k, v));
	}

private:
	std::map<Literal, int> lookup_;
};

namespace {
	void fatal_unresolved_literal_error(ConstantAST const& v)
	{
		Literal const* lit = boost::get<Literal>(&v);
		assert(lit);

		std::cout << "Unable to resolve the value of literal '" << std::string(lit->begin(), lit->end()) << "'" << std::endl;
		::abort();
	}
}

Program process(ProgramAST const& program)
{
	std::map<Literal, ConstantAST> unresolved;
	ConstResolver resolver;

	BOOST_FOREACH(StatementAST const& stmt, program) {
		if (DeclarationAST const* decl = boost::get<DeclarationAST>(&stmt)) {
			unresolved.insert(std::make_pair(decl->dname, decl->dvalue));
		} else if (InstructionAST const* inst = boost::get<InstructionAST>(&stmt)) {

			if (DeclarationAST const* d = boost::get<DeclarationAST>(&inst->id))
				unresolved.insert(std::make_pair(d->dname, d->dvalue));
			if (DeclarationAST const* d = boost::get<DeclarationAST>(&inst->op))
				unresolved.insert(std::make_pair(d->dname, d->dvalue));
			if (DeclarationAST const* d = boost::get<DeclarationAST>(&inst->argA))
				unresolved.insert(std::make_pair(d->dname, d->dvalue));
			if (DeclarationAST const* d = boost::get<DeclarationAST>(&inst->argB))
				unresolved.insert(std::make_pair(d->dname, d->dvalue));
			if (DeclarationAST const* d = boost::get<DeclarationAST>(&inst->ret))
				unresolved.insert(std::make_pair(d->dname, d->dvalue));
		} else {
			assert(false);
		}
	}


	BOOST_FOREACH(Mnemonic mnem, all_mnemonics()) {
		std::string as_string(show(mnem));

		if (!has_return(mnem))
			resolver.add_constant(as_string, to_opcode(mnem, CHANGE_ARGA)); // TODO: <- should only be for mnemonics that don't return anything

		resolver.add_constant(as_string + "o", to_opcode(mnem, CHANGE_OP));
		resolver.add_constant(as_string + "a", to_opcode(mnem, CHANGE_ARGA));
		resolver.add_constant(as_string + "b", to_opcode(mnem, CHANGE_ARGB));
		resolver.add_constant(as_string + "r", to_opcode(mnem, CHANGE_RET));

	}

	const int max_passes = 10;
	for (int i = 0 ; !unresolved.empty() ; ++i) {
		if (i == max_passes) {
			std::cerr << "Could not figure out all constants (in << " << max_passes << " passes). Circular dep?";
			Literal l(unresolved.begin()->first);
			std::cerr << "First constant not resolved was: " << std::string(l.begin(), l.end()) << std::endl;
			::abort();
		}

		for (std::map<Literal, ConstantAST>::iterator it(unresolved.begin()) ; it != unresolved.end() ; )
		{
			boost::optional<int> val(boost::apply_visitor(resolver, it->second));

			if (val.is_initialized()) {
				resolver.add_constant(it->first, *val);
				unresolved.erase(it++);
			} else {
				++it;
			}
		}
	}

	Program prog;

	// Now that we know all the constants, it's all pretty simple.
	BOOST_FOREACH(StatementAST const& stmt, program) {

		if (InstructionAST const* inst = boost::get<InstructionAST>(&stmt)) {

			boost::optional<int> id = boost::apply_visitor(resolver, inst->id);
			boost::optional<int> op = boost::apply_visitor(resolver, inst->op);
			boost::optional<int> argA = boost::apply_visitor(resolver, inst->argA);
			boost::optional<int> argB = boost::apply_visitor(resolver, inst->argB);
			boost::optional<int> ret = boost::apply_visitor(resolver, inst->ret);

			if (!id.is_initialized())
				fatal_unresolved_literal_error(inst->id);
			if (!op.is_initialized())
				fatal_unresolved_literal_error(inst->op);
			if (!argA.is_initialized())
				fatal_unresolved_literal_error(inst->argA);
			if (!argB.is_initialized())
				fatal_unresolved_literal_error(inst->argB);
			if (!ret.is_initialized())
				fatal_unresolved_literal_error(inst->ret);


			prog.push_back(Instruction(*id, *op, *argA, *argB, *ret));

		} else {
			assert( boost::get<DeclarationAST>(&stmt) );
		}
	}

	return prog;

}


Program assemble(std::string const& contents)
{
	assert(!contents.empty());

	std::string::const_iterator fwd_begin(contents.begin());
	std::string::const_iterator fwd_end(contents.end());

	ProgramGrammar<std::string::const_iterator> grammar;
	ProgramAST outputAST;
	qi::phrase_parse(fwd_begin, fwd_end, grammar, ( (qi::space - qi::eol) |  (qi::lit('#') >> *(qi::char_ - qi::eol)) ), outputAST);


	if (fwd_begin != fwd_end) {
		std::cout << "Parsing failed!! Printing characters from immediately after parsing fail: " << std::endl;
		const int max_chars = 30;
		int char_counter(0);
		while (fwd_begin != fwd_end && char_counter < max_chars) {
			std::cout << *fwd_begin;
			++char_counter;
			++fwd_begin;
		}
		::exit(EXIT_FAILURE);
	}


	return process(outputAST);
}

Program assemble(std::istream & instream)
{
	assert(!instream.fail());
	instream.unsetf(std::ios_base::skipws);

	std::string contents(std::istreambuf_iterator<char>(instream.rdbuf() ), std::istreambuf_iterator<char>());

	return assemble(contents);
}

